import os
import img2pdf
from PIL import Image
import uuid
import io
from fastapi import APIRouter, UploadFile, File, HTTPException, Query
from typing import List, Optional

# 创建路由器
router = APIRouter(
    prefix="/image2pdf",
    tags=["图片转PDF"],
    responses={404: {"description": "Not found"}},
)

# 默认目录
DEFAULT_INPUT_DIR = "input_files"
DEFAULT_OUTPUT_DIR = "output_files"

def ensure_directory_exists(directory_path):
    """确保目录存在，如果不存在则创建"""
    # 检查路径是否为空字符串或None
    if not directory_path:
        return  # 如果是空字符串或None，直接返回，不创建目录
        
    if not os.path.exists(directory_path):
        os.makedirs(directory_path)

def is_image_file(file_path):
    """检查文件是否为图片"""
    try:
        img = Image.open(file_path)
        img.verify()  # 验证图片完整性
        return True
    except Exception:
        return False

def get_image_dpi(image_path):
    """获取图片的DPI信息，如果不存在则返回默认值(96,96)"""
    try:
        with Image.open(image_path) as img:
            dpi = img.info.get('dpi', (96, 96))
            return dpi
    except Exception:
        return (96, 96)

def process_image_to_pdf(image_path, output_dir="output_files"):
    """
    将单张图片转换为PDF文件，使用img2pdf库以无损方式转换
    
    Args:
        image_path: 图片文件路径
        output_dir: 输出目录，默认为'output_files'
        
    Returns:
        生成的PDF文件路径
    """
    # 确保输出目录存在
    ensure_directory_exists(output_dir)
    
    # 检查文件是否为图片
    if not is_image_file(image_path):
        raise ValueError(f"文件不是有效的图片: {image_path}")
    
    # 生成唯一的输出文件名
    image_basename = os.path.basename(image_path)
    output_filename = f"output_image_{os.path.splitext(image_basename)[0]}_{uuid.uuid4().hex[:8]}.pdf"
    pdf_path = os.path.join(output_dir, output_filename)
    
    # 使用img2pdf将图片转换为PDF（无损转换）
    with open(pdf_path, "wb") as f:
        try:
            # 尝试直接转换
            f.write(img2pdf.convert(image_path))
        except Exception as e:
            # 如果直接转换失败（例如CMYK图像等），尝试通过PIL处理
            print(f"直接转换失败，尝试通过PIL处理: {str(e)}")
            img = Image.open(image_path)
            
            # 如果图片有透明通道但不是RGBA，转换为RGBA
            if img.mode == 'P' and 'transparency' in img.info:
                img = img.convert('RGBA')
                
            # 如果是RGBA，需要去除透明度（img2pdf不支持透明度）
            if img.mode == 'RGBA':
                background = Image.new('RGB', img.size, (255, 255, 255))
                background.paste(img, mask=img.split()[3])  # 3是alpha通道
                img = background
            
            # 如果是CMYK，转换为RGB
            elif img.mode == 'CMYK':
                img = img.convert('RGB')
                
            # 保存为临时BytesIO对象
            with io.BytesIO() as output:
                img.save(output, format="PNG")
                output.seek(0)
                f.write(img2pdf.convert(output))
    
    return pdf_path

def batch_process_images_to_pdf(image_paths, output_dir="output_files"):
    """
    批量将多张图片转换为单个PDF文件，使用img2pdf库以无损方式转换
    
    Args:
        image_paths: 图片文件路径列表
        output_dir: 输出目录，默认为'output_files'
        
    Returns:
        生成的PDF文件路径
    """
    # 确保输出目录存在
    ensure_directory_exists(output_dir)
    
    # 生成唯一的输出文件名
    output_filename = f"output_batch_images_{uuid.uuid4().hex[:8]}.pdf"
    pdf_path = os.path.join(output_dir, output_filename)
    
    # 筛选有效的图片文件
    valid_image_paths = []
    processed_images = []
    
    for image_path in image_paths:
        # 检查文件是否为图片
        if not is_image_file(image_path):
            continue
            
        try:
            img = Image.open(image_path)
            
            # 处理特殊格式的图片
            if img.mode == 'P' and 'transparency' in img.info:
                img = img.convert('RGBA')
                
            if img.mode == 'RGBA':
                # 处理透明背景
                background = Image.new('RGB', img.size, (255, 255, 255))
                background.paste(img, mask=img.split()[3])
                img = background
            elif img.mode == 'CMYK':
                # 处理CMYK颜色空间
                img = img.convert('RGB')
                
            # 如果图片需要处理，保存为临时文件
            if img.mode != Image.open(image_path).mode:
                temp_path = os.path.join(output_dir, f"temp_{uuid.uuid4().hex[:8]}.png")
                img.save(temp_path)
                processed_images.append(temp_path)
                valid_image_paths.append(temp_path)
            else:
                valid_image_paths.append(image_path)
                
        except Exception as e:
            print(f"处理图片 {image_path} 时出错: {str(e)}")
            continue
    
    if not valid_image_paths:
        raise ValueError("没有有效的图片文件可以转换")
    
    try:
        # 使用img2pdf将所有图片合并为一个PDF
        with open(pdf_path, "wb") as f:
            # 获取布局函数（保持图片原始尺寸）
            layout_fun = img2pdf.get_layout_fun(None)
            f.write(img2pdf.convert(valid_image_paths, layout_fun=layout_fun))
        
        # 清理临时文件
        for temp_file in processed_images:
            if os.path.exists(temp_file):
                os.remove(temp_file)
                
        return pdf_path
    except Exception as e:
        # 如果转换失败，也要清理临时文件
        for temp_file in processed_images:
            if os.path.exists(temp_file):
                os.remove(temp_file)
        raise ValueError(f"生成PDF文件失败: {str(e)}")

# 添加路由器端点
@router.post("/convert")
async def convert_images_to_pdf(
    files: List[UploadFile] = File(...),
    input_dir: Optional[str] = Query(DEFAULT_INPUT_DIR, description="输入文件目录"),
    output_dir: Optional[str] = Query(DEFAULT_OUTPUT_DIR, description="输出文件目录")
):
    """
    将多张图片转换为单个PDF文件（使用img2pdf库实现无损转换）
    
    - **files**: 要上传的图片文件列表
    - **input_dir**: 输入文件保存目录，默认为'input_files'
    - **output_dir**: 输出文件保存目录，默认为'output_files'
    
    注意：此接口使用img2pdf库实现无损转换，保持原始图像质量，比传统方法文件更小，转换更快。
    支持的图片格式：JPEG、PNG、TIFF等，详情参考：https://github.com/josch/img2pdf
    """
    #
    # 处理默认目录
    if not input_dir or input_dir.strip() == "":
        input_dir = DEFAULT_INPUT_DIR
    
    if not output_dir or output_dir.strip() == "":
        output_dir = DEFAULT_OUTPUT_DIR
    
    # 确保目录存在
    ensure_directory_exists(input_dir)
    ensure_directory_exists(output_dir)
    
    # 检查是否有文件上传
    if not files or len(files) == 0:
        raise HTTPException(status_code=400, detail="请至少上传一个图片文件")
    
    # 保存所有上传的文件
    file_paths = []
    valid_extensions = ['.jpg', '.jpeg', '.png', '.gif', '.bmp', '.webp', '.tiff', '.tif']
    
    for file in files:
        # 检查文件类型
        file_ext = os.path.splitext(file.filename.lower())[1]
        if file_ext not in valid_extensions:
            continue  # 跳过不支持的文件类型
        
        # 保存文件
        file_path = os.path.join(input_dir, file.filename)
        with open(file_path, "wb") as buffer:
            buffer.write(await file.read())
        file_paths.append(file_path)
    
    if not file_paths:
        raise HTTPException(status_code=400, detail=f"没有有效的图片文件。支持的格式: {', '.join(valid_extensions)}")
    
    try:
        # 处理图片文件
        if len(file_paths) == 1:
            # 单张图片
            output_path = process_image_to_pdf(file_paths[0], output_dir=output_dir)
        else:
            # 多张图片合并为一个PDF
            output_path = batch_process_images_to_pdf(file_paths, output_dir=output_dir)
        
        # 返回文件下载URL
        file_name = os.path.basename(output_path)
        download_url = f"/outputs/{file_name}"
        
        return {
            "success": True,
            "message": f"共处理 {len(file_paths)} 张图片，生成无损PDF成功",
            "file_path": output_path,
            "download_url": download_url
        }
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"处理文件时出错: {str(e)}")
